from collections.abc import Iterable, Iterator
from typing import Any, Callable, Generic, TypeVar

from zkl_pyutils_fsspec import FsLike, resolve_child_fs

from zkl_aiutils_datasets.basics import Dataset, DatasetWriter
from zkl_aiutils_datasets.loading import load_dataset_by_scripting_from_itself
from zkl_aiutils_datasets.processing import ChainedDatasetsWriter, DelegatingDataset, DelegatingDatasetWriter, \
    KnownSizeChainedDatasets
from zkl_aiutils_datasets.scripting import DatasetMetadata, make_dataset_metadata, write_dataset_metadata
from .datasets import DirectoryDatasets

AnySample = TypeVar('AnySample', bound=Any)
AnyDataset = TypeVar('AnyDataset', bound=Dataset)
AnyDatasetWriter = TypeVar('AnyDatasetWriter', bound=DatasetWriter)


class ChainedDirectoryDatasets(
    DelegatingDataset[AnySample],
    DirectoryDatasets[AnyDataset],
    Generic[AnySample, AnyDataset],
):
    def __init__(self,
        fs: FsLike,
        child_factory: Callable[[FsLike], AnyDataset] = load_dataset_by_scripting_from_itself,
    ):
        DirectoryDatasets.__init__(self, fs, child_factory)

        children_samples_n = tuple(
            child_metadata.get('samples_n', None)
            for child_metadata in self.children_metadata.values()
        ) if self.children_metadata is not None else None
        delegated_dataset = KnownSizeChainedDatasets(self.indexed_children, children_samples_n)
        DelegatingDataset.__init__(self, delegated_dataset)

    def __len__(self) -> int:
        if self.metadata is not None:
            samples_n = self.metadata.get('samples_n', None)
            if samples_n is not None:
                return samples_n
        return super().__len__()


class ChainedDirectoryDatasetsWriter(
    ChainedDatasetsWriter[AnySample],
    Generic[AnySample, AnyDatasetWriter],
):
    def __init__(self,
        fs: FsLike,
        names: Iterable[str] | Iterator[str],
        child_factory: Callable[[FsLike], AnyDatasetWriter], *,
        loader: Callable | type | None = ChainedDirectoryDatasets,
    ):
        self._fs = fs
        self._metadata = make_dataset_metadata(loader)

        children = self._iter_children(names, child_factory)
        ChainedDatasetsWriter.__init__(self, children)

        if self._metadata:
            write_dataset_metadata(self._fs, self._metadata)

    def _iter_children(self,
        names: Iterable[str] | Iterator[str],
        child_factory: Callable[[FsLike], AnyDatasetWriter],
    ) -> Iterator[DatasetWriter[AnySample]]:
        for name in names:
            if name in self._metadata.get('children', {}):
                raise ValueError(f'Found duplicate child name {name!r}!')
            child_writer = child_factory(resolve_child_fs(self._fs, name))
            yield ChainedDirectoryDatasetsChildWriter(self, child_writer, name)

    def _submit_child_metadata(self, child_name: str, child_metadata: DatasetMetadata | Any):
        self._metadata['children'] = {
            **self._metadata.get('children', {}),
            child_name: child_metadata}
        self._metadata['samples_n'] = sum(
            child_metadata.get('samples_n', 0)
            for child_metadata in self._metadata.get('children', {}).values())
        write_dataset_metadata(self._fs, self._metadata)


class ChainedDirectoryDatasetsChildWriter(DelegatingDatasetWriter[AnySample]):
    def __init__(self,
        parent: ChainedDirectoryDatasetsWriter[AnySample, AnyDatasetWriter],
        child: DatasetWriter[AnySample],
        name: str,
    ):
        self._parent = parent
        super().__init__(child)

        self._name = name
        self._metadata = DatasetMetadata()
        self._metadata_dirty = False

    def _submit_metadata(self):
        if self._metadata_dirty:
            self._metadata_dirty = False
            self._parent._submit_child_metadata(self._name, self._metadata)

    def write(self, sample: AnySample, /):
        super().write(sample)
        self._metadata_dirty = True
        self._metadata['samples_n'] = self._metadata.get('samples_n', 0) + 1

    def flush(self):
        super().flush()
        self._submit_metadata()

    def close(self):
        super().close()
        self._submit_metadata()
